#!/usr/bin/env python
# -*- conding: utf-8 -*-

"""
@Time     : 2024/9/7 22:25
@Author   : liujingmao
@File     : 1.带阈值的相似性搜索.py
"""

import dotenv
from langchain_community.vectorstores import FAISS
from langchain_core.documents import Document
from langchain_openai import OpenAIEmbeddings

dotenv.load_dotenv()

# 1. create embedding OpenAIEmbeddings 实例
embedding = OpenAIEmbeddings(model="text-embedding-3-small")

# 2. 创建文件documents
documents = [
    Document(page_content="笨笨是一只很喜欢睡觉的狗", metadata={"page": 1}),
    Document(page_content="我喜欢在夜晚听音乐，这让我感到放松。", metadata={"page": 2}),
    Document(page_content="猫咪在窗台上打盹，看起来非常可爱。", metadata={"page": 3}),
    Document(page_content="学习新技能是每个人都应该追求的目标。", metadata={"page": 4}),
    Document(page_content="我最喜欢的食物是意大利面，尤其是番茄酱的那种。", metadata={"page": 5}),
    Document(page_content="昨晚我做了一个奇怪的梦，梦见自己在太空飞行。", metadata={"page": 6}),
    Document(page_content="我的手机突然关机了，让我有些焦虑。", metadata={"page": 7}),
    Document(page_content="阅读是我每天都会做的事情，我觉得很充实。", metadata={"page": 8}),
    Document(page_content="他们一起计划了一次周末的野餐，希望天气能好。", metadata={"page": 9}),
    Document(page_content="我的狗名字叫笨笨，喜欢追逐球，看起来非常开心。", metadata={"page": 10}),
]

# 3. 将documents内容传递给大语言模型，并生成一个db实例
db = FAISS.from_documents(documents, embedding)

# 4. 通过传递的query和score_threshold在db实例中找到相关的内容，并打印出来
print(db.similarity_search_with_relevance_scores("我养了一条狗，它的名字叫笨笨", score_threshold=0.4))
